# -*- coding: utf-8 -*-

import os
from transformers import AutoTokenizer, AutoModel

# 设置模型的保存目录
save_directory = "/home/ubuntu/code/git/subject-word-extraction/chatglm3_model"  # 替换为您想要的目录路径
os.makedirs(save_directory, exist_ok=True)

# 指定模型名称
model_name = "THUDM/chatglm3-6b"

# 下载并保存 tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=save_directory, trust_remote_code=True)

# 下载并保存模型
model = AutoModel.from_pretrained(model_name, cache_dir=save_directory, trust_remote_code=True)

# 检查模型和 tokenizer 是否下载完成
print(f"Model and tokenizer have been downloaded to {save_directory}")
    